import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer
from model import BertSpamClassifier
from sklearn.metrics import classification_report
from cot_utils import generate_cot

PRETRAINED_MODEL = 'bert-base-uncased'
BATCH_SIZE = 32
MAX_LEN = 128
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 1. 加载数据
df = pd.read_csv('data/sms_spam_cot.csv')
# 只用10%做测试
_, test_df = df[:int(len(df)*0.9)], df[int(len(df)*0.9):]

# 2. 数据集定义
class SpamDataset(Dataset):
    def __init__(self, df, tokenizer, max_len):
        self.texts = (df['text'] + ' [SEP] ' + df['cot']).tolist()
        self.labels = df['label'].map({'ham': 0, 'spam': 1}).tolist()
        self.tokenizer = tokenizer
        self.max_len = max_len
    def __len__(self):
        return len(self.texts)
    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            truncation=True,
            padding='max_length',
            max_length=self.max_len,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL)
test_dataset = SpamDataset(test_df, tokenizer, MAX_LEN)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)

# 3. 加载模型
model = BertSpamClassifier(PRETRAINED_MODEL)
model.load_state_dict(torch.load('spam_bert_cot.pth', map_location=DEVICE))
model.to(DEVICE)
model.eval()

# 4. 批量评估
all_preds, all_labels = [], []
with torch.no_grad():
    for batch in test_loader:
        input_ids = batch['input_ids'].to(DEVICE)
        attention_mask = batch['attention_mask'].to(DEVICE)
        labels = batch['label'].to(DEVICE)
        logits = model(input_ids, attention_mask)
        preds = torch.argmax(logits, dim=1).cpu().numpy()
        all_preds.extend(preds)
        all_labels.extend(labels.cpu().numpy())

print(classification_report(all_labels, all_preds, target_names=['ham', 'spam'])) 